import numpy as np
import matplotlib.pyplot as plt
from math import cos, sin, pi, sqrt, atan2, hypot

# 定义路径段类型常量
L_SEG = 1  # 左转前进
S_SEG = 2  # 直行
R_SEG = 3  # 右转前进
L_SEG_B = -1  # 左转后退
R_SEG_B = -3  # 右转后退


# 路径段结构
class PathSegment:
    def __init__(self, seg_type, length):
        self.type = seg_type  # 段类型
        self.length = length  # 段长度
        self.radius = 0.0  # 转弯半径（仅用于圆弧段）


# Reeds-Shepp路径类
class ReedsSheppPath:
    def __init__(self):
        self.segments = []  # 路径段列表
        self.total_length = 0.0  # 路径总长度


# 主路径规划函数
def reeds_shepp_path(start, goal, radius, step_size=0.1):
    """
    计算Reeds-Shepp路径
    :param start: 起始位姿 (x, y, yaw)
    :param goal: 目标位姿 (x, y, yaw)
    :param radius: 最小转弯半径
    :param step_size: 路径点采样步长
    :return: 路径点列表 [(x, y, yaw)]
    """
    # 转换为标准问题（起点在原点，方向为0）
    dx   = goal[0] - start[0]
    dy   = goal[1] - start[1]
    dyaw = goal[2] - start[2]

    # 转换为车辆坐标系
    x = (cos(start[2]) * dx + sin(start[2]) * dy)
    y = (-sin(start[2]) * dx + cos(start[2]) * dy)
    yaw = dyaw

    # 计算所有可能路径并选择最短的
    paths = []
    for path_type in range(1, 49):
        path = calc_path(x, y, yaw, radius, path_type)
        if path:
            paths.append(path)

    if not paths:
        return None

    # 选择最短路径
    best_path = min(paths, key=lambda p: p.total_length)

    # 生成路径点
    return generate_path_points(start, best_path, radius, step_size)


def calc_path(x, y, yaw, radius, path_type):
    """
    计算特定类型的路径
    :return: ReedsSheppPath对象或None（若无效）
    """
    path = ReedsSheppPath()

    # 根据路径类型配置段序列
    if path_type == 1:  # CSC类型: LSL
        seg_types = [L_SEG, S_SEG, L_SEG]
        lengths = calc_lsl(x, y, yaw, radius)
    elif path_type == 2:  # LSR
        seg_types = [L_SEG, S_SEG, R_SEG]
        lengths = calc_lsr(x, y, yaw, radius)
    # 其他45种类型需要类似实现（此处省略）
    else:
        return None

    if lengths is None:
        return None

    # 构建路径段
    for seg_type, length in zip(seg_types, lengths):
        if abs(length) < 1e-6:  # 忽略零长度段
            continue
        segment = PathSegment(seg_type, abs(length))
        segment.radius = radius if abs(seg_type) in [1, 3] else float('inf')
        path.segments.append(segment)
        path.total_length += abs(length)

    return path


def calc_lsl(x, y, yaw, radius):
    """计算LSL路径段长度"""
    # 计算中间点
    u = x - radius * sin(yaw) + radius * sin(0)
    v = y + radius * cos(yaw) - radius * cos(0)
    t = atan2(v, u)
    t = t if t >= 0 else t + 2 * pi

    # 计算各段长度
    seg1 = t
    seg2 = hypot(u, v)
    seg3 = yaw - t

    return [seg1 * radius, seg2, seg3 * radius]


def calc_lsr(x, y, yaw, radius):
    """计算LSR路径段长度"""
    # 计算中间点
    u = x + radius * sin(yaw) - radius * sin(0)
    v = y - radius * cos(yaw) + radius * cos(0)
    t1 = atan2(v, u)
    t1 = t1 if t1 >= 0 else t1 + 2 * pi

    u2 = u - 2 * radius * sin(t1)
    v2 = v + 2 * radius * cos(t1)
    t2 = atan2(v2, u2) if hypot(u2, v2) > 1e-6 else 0

    # 计算各段长度
    seg1 = t1
    seg2 = hypot(u2, v2)
    seg3 = t2 - t1 + yaw

    return [seg1 * radius, seg2, seg3 * radius]


def generate_path_points(start, path, radius, step_size):
    """
    生成路径点
    :param start: 起始位姿 (x, y, yaw)
    :param path: ReedsSheppPath对象
    :param radius: 转弯半径
    :param step_size: 采样步长
    :return: 路径点列表
    """
    points = []
    x, y, yaw = start

    # 添加起点
    points.append((x, y, yaw))

    # 遍历所有段
    for seg in path.segments:
        length = seg.length
        steps = int(length / step_size)

        for i in range(1, steps + 1):
            s = i * step_size
            if seg.type == S_SEG:  # 直线段
                dx = s * cos(yaw)
                dy = s * sin(yaw)
                dyaw = 0
            else:  # 圆弧段
                # 计算曲率中心
                direction = 1 if seg.type > 0 else -1
                turn_dir = -1 if seg.type in [L_SEG, L_SEG_B] else 1
                center_x = x - direction * radius * sin(yaw) * turn_dir
                center_y = y + direction * radius * cos(yaw) * turn_dir

                # 计算圆弧上的点
                angle = direction * s / radius
                dx = center_x + radius * sin(yaw + angle) * turn_dir
                dy = center_y - radius * cos(yaw + angle) * turn_dir
                dyaw = angle

            # 更新位姿
            x = dx
            y = dy
            yaw = (yaw + dyaw) % (2 * pi)
            points.append((x, y, yaw))

    return points


# 示例使用
if __name__ == "__main__":
    # 参数设置
    start_pose = (0, 0, 0)  # 起始位姿 (x, y, yaw)
    goal_pose = (4, 4, -pi / 2)  # 目标位姿
    min_radius = 5.0  # 最小转弯半径
    step = 0.1  # 路径点步长

    # 计算路径
    path_points = reeds_shepp_path(start_pose, goal_pose, min_radius, step)

    if path_points:
        # 可视化路径
        x_vals = [p[0] for p in path_points]
        y_vals = [p[1] for p in path_points]

        plt.figure(figsize=(20, 16))
        plt.plot(x_vals, y_vals, 'b-', label='Path')
        plt.quiver(start_pose[0], start_pose[1],
                   cos(start_pose[2]), sin(start_pose[2]),
                   scale=10, color='r', label='Start')
        plt.quiver(goal_pose[0], goal_pose[1],
                   cos(goal_pose[2]), sin(goal_pose[2]),
                   scale=10, color='g', label='Goal')
        plt.axis('equal')
        plt.grid(True)
        plt.legend()
        plt.title('Reeds-Shepp Path with Turning Radius Constraint')
        plt.show()
    else:
        print("No valid path found!")